import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiTaskLoss(nn.Module):
    def __init__(self, intent_weight=1.0, trajectory_weight=1.0, reg_weight=0.01):
        super().__init__()
        self.intent_weight = intent_weight
        self.trajectory_weight = trajectory_weight
        self.reg_weight = reg_weight
        
        # 意图预测使用Focal Loss处理类别不平衡
        self.intent_loss = FocalLoss()
        
        # 轨迹预测使用平滑L1损失
        self.trajectory_loss = nn.SmoothL1Loss()
    
    def forward(self, intent_pred, trajectory_pred, labels):
        intent_gt = labels['intent']
        trajectory_gt = labels['trajectory']
        
        # 意图损失
        intent_loss = self.intent_loss(intent_pred, intent_gt)
        
        # 轨迹损失
        trajectory_loss = self.trajectory_loss(trajectory_pred, trajectory_gt)
        
        # 总损失
        total_loss = (self.intent_weight * intent_loss + 
                     self.trajectory_weight * trajectory_loss)
        
        return total_loss

class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction
    
    def forward(self, inputs, targets):
        BCE_loss = F.binary_cross_entropy(inputs, targets, reduction='none')
        pt = torch.exp(-BCE_loss)
        F_loss = self.alpha * (1-pt)**self.gamma * BCE_loss
        
        if self.reduction == 'mean':
            return torch.mean(F_loss)
        elif self.reduction == 'sum':
            return torch.sum(F_loss)
        else:
            return F_loss